/*
 * Copyright © 2024-2025 Wenze Wei
 *
 * This file is part of StadionOS.
 *
 * This Source Code Form is subject to the terms of the GNU General Public License, v. 2.0.
 * You can obtain a copy of the GPL at
 * https://www.gnu.org/licenses/old-licenses/gpl-2.0.html.
 *
 * Software distributed under the License is distributed on an "AS IS" basis,
 * WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License for
 * the specific language governing rights and limitations under the License.
 */

#include "compositor.h"
#include "output.h"
#include "animation.h"

static void handle_new_output(struct wl_listener *listener, void *data);
static void handle_new_input(struct wl_listener *listener, void *data);
static void handle_new_xdg_surface(struct wl_listener *listener, void *data);
static void handle_new_xwayland_surface(struct wl_listener *listener, void *data);
static void handle_window_destroy(struct wl_listener *listener, void *data);
static void handle_window_map(struct wl_listener *listener, void *data);
static void handle_window_commit(struct wl_listener *listener, void *data);
static bool socket_exists(const char *path);
static char *generate_unique_socket_name(const char *base_name);

static bool socket_exists(const char *path) {
    return access(path, F_OK) == 0;
}

static char *generate_unique_socket_name(const char *base_name) {
    const char *runtime_dir = getenv("XDG_RUNTIME_DIR");
    if (!runtime_dir) {
        runtime_dir = "/run/user";
        char user_runtime_dir[256];
        snprintf(user_runtime_dir, sizeof(user_runtime_dir), "%s/%d", runtime_dir, getuid());

        if (access(user_runtime_dir, F_OK) == 0) {
            runtime_dir = user_runtime_dir;
            fprintf(stderr, "XDG_RUNTIME_DIR is not set, using fallback directory: %s\n", runtime_dir);
        } else {
            runtime_dir = "/tmp";
            fprintf(stderr, "XDG_RUNTIME_DIR is not set and /run/user/%d does not exist, using fallback directory: %s\n", getuid(), runtime_dir);
        }
    }

    char *socket_path = malloc(512);
    if (!socket_path) {
        fprintf(stderr, "Failed to allocate memory for socket path\n");
        return NULL;
    }

    int index = 0;
    char *prev_path = NULL;

    while (1) {
        if (prev_path) {
            free(prev_path);
        }
        prev_path = strdup(socket_path);
        
        if (index == 0) {
            snprintf(socket_path, 512, "%s/%s", runtime_dir, base_name);
        } else {
            snprintf(socket_path, 512, "%s/%s-%d", runtime_dir, base_name, index);
        }

        if (!socket_exists(socket_path)) {
            free(prev_path);
            return socket_path;
        }

        index++;
    }
}

struct onatu_compositor *compositor_get_instance() {
    return &global_compositor;
}

void compositor_init(struct onatu_compositor *compositor) {
    compositor->xwayland = wlr_xwayland_create(compositor->display, compositor->xwm, true);
    wl_signal_add(&compositor->xwayland->events.new_surface, &compositor->new_xwayland_surface);

    if (wlr_renderer_get_drm_format(compositor->renderer, DRM_FORMAT_MOD_LINEAR) == 0) {
        LOG(compositor, "Hardware acceleration unavailable, falling back to software");
        compositor->software_fallback = true;
    }

    compositor->log_monitor = wl_event_loop_add_fd(wl_display_get_event_loop(compositor->display),
        fileno(compositor->log_file), WL_EVENT_READABLE, handle_log_event, compositor);

    struct wlr_xdg_decoration_manager_v1 *decoration_manager = 
        wlr_xdg_decoration_manager_v1_create(compositor->display);
    
    wl_signal_add(&decoration_manager->events.new_toplevel_decoration,
        &compositor->new_decoration);
    compositor->log_file = fopen("/var/log/onatu.log", "a");
    if (!compositor->log_file) {
        fprintf(stderr, "Failed to open log file");
        exit(1);
    }

    compositor->scene = wlr_scene_create();
    compositor->scene_layout = wlr_scene_output_layout_create(compositor->scene);

    compositor->display = wl_display_create();
    if (!compositor->display) {
        LOG(compositor, "FATAL: Failed to create Wayland display");
        exit(1);
    }

    compositor->backend = wlr_multi_backend_create(compositor->display);
    struct wlr_backend *libinput = wlr_libinput_backend_create(compositor->display, compositor->session);
    struct wlr_backend *drm = wlr_drm_backend_create(compositor->display, compositor->session);
    wlr_multi_backend_add(compositor->backend, libinput);
    wlr_multi_backend_add(compositor->backend, drm);
    if (!compositor->backend) {
        LOG(compositor, "FATAL: Failed to create backend");
        compositor_destroy(compositor);
        exit(1);
    }

    compositor->renderer = wlr_backend_get_renderer(compositor->backend);
    
    EGLint config_attrs[] = {
        EGL_SURFACE_TYPE, EGL_WINDOW_BIT,
        EGL_RED_SIZE, 8,
        EGL_GREEN_SIZE, 8,
        EGL_BLUE_SIZE, 8,
        EGL_ALPHA_SIZE, 8,
        EGL_RENDERABLE_TYPE, EGL_OPENGL_ES2_BIT,
        EGL_NONE
    };

    EGLConfig config;
    EGLint num_config;
    EGLDisplay egl_display = wlr_gles2_renderer_get_egl_display(compositor->renderer);
    eglChooseConfig(egl_display, config_attrs, &config, 1, &num_config);
    
    EGLint ctx_attrs[] = {
        EGL_CONTEXT_CLIENT_VERSION, 2,
        EGL_NONE
    };
    EGLContext egl_ctx = eglCreateContext(egl_display, config, EGL_NO_CONTEXT, ctx_attrs);
    
    wlr_renderer_init_wl_display(compositor->renderer);

    compositor->allocator = wlr_allocator_autocreate(compositor->backend,
        compositor->renderer);
    
    if (!compositor->renderer || !compositor->allocator) {
        fprintf(stderr, "Failed to create renderer or allocator\n");
        compositor_destroy(compositor);
        exit(1);
    }

    if (!wlr_renderer_init_wl_display(compositor->renderer, compositor->display)) {
        fprintf(stderr, "Failed to initialize renderer with Wayland display\n");
        exit(1);
    }

    wl_list_init(&compositor->keyboards);
    wl_list_init(&compositor->pointers);
    wl_list_init(&compositor->outputs);
    wl_list_init(&compositor->windows);

    compositor->new_output.notify = handle_new_output;
    wl_signal_add(&compositor->backend->events.new_output, &compositor->new_output);

    compositor->new_input.notify = handle_new_input;
    wl_signal_add(&compositor->backend->events.new_input, &compositor->new_input);

    compositor->xdg_shell = wlr_xdg_shell_create(compositor->display, XDG_SHELL_VERSION_STABLE);
    
    compositor->border_rect = wlr_scene_rect_create(&compositor->scene->tree, 8, 8, (float[4]){0.3, 0.5, 0.8, 0.5});
    wlr_scene_node_set_enabled(&compositor->border_rect->node, false);
    compositor->new_xdg_surface.notify = handle_new_xdg_surface;
    wl_signal_add(&compositor->xdg_shell->events.new_surface, &compositor->new_xdg_surface);

    compositor->seat = wlr_seat_create(compositor->display, "seat0");
}

int compositor_run(struct onatu_compositor *compositor) {
    if (!wlr_backend_start(compositor->backend)) {
        fprintf(stderr, "Failed to start backend\n");
        return 1;
    }

    char *socket_path __attribute__((cleanup(free))) = generate_unique_socket_name("Onatu");
    if (!socket_path) {
        fprintf(stderr, "Failed to generate unique socket name\n");
        return 1;
    }

    const char *socket_name = strrchr(socket_path, '/') + 1;
    if (setenv("WAYLAND_DISPLAY", socket_name, 1) != 0) {
        fprintf(stderr, "Failed to set WAYLAND_DISPLAY\n");
        free(socket_path);
        return 1;
    }

    unlink(socket_path);

    int sockfd = socket(AF_UNIX, SOCK_STREAM, 0);
    if (sockfd < 0) {
        fprintf(stderr, "Failed to create socket: %s\n", strerror(errno));
        free(socket_path);
        return 1;
    }

    struct sockaddr_un addr;
    memset(&addr, 0, sizeof(addr));
    addr.sun_family = AF_UNIX;
    strncpy(addr.sun_path, socket_path, sizeof(addr.sun_path) - 1);

    if (bind(sockfd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        fprintf(stderr, "Failed to bind socket: %s\n", strerror(errno));
        close(sockfd);
        free(socket_path);
        return 1;
    }

    if (listen(sockfd, 5) < 0) {
        fprintf(stderr, "Failed to listen on socket: %s\n", strerror(errno));
        close(sockfd);
        free(socket_path);
        return 1;
    }

    if (wl_display_add_socket_fd(compositor->display, sockfd) < 0) {
        fprintf(stderr, "Failed to add socket to Wayland display\n");
        close(sockfd);
        free(socket_path);
        return 1;
    }

    printf("Onatu is running on WAYLAND_DISPLAY=%s\n", socket_name);
    wl_display_run(compositor->display);

    close(sockfd);
    free(socket_path);
    return 0;
}

void compositor_destroy(struct onatu_compositor *compositor) {
    if (compositor->backend) {
        wlr_backend_destroy(compositor->backend);
    }
    if (compositor->renderer) {
        wlr_renderer_destroy(compositor->renderer);
    }
    if (compositor->allocator) {
        wlr_allocator_destroy(compositor->allocator);
    }
    if (compositor->scene) {
        wlr_scene_destroy(compositor->scene);
    }
    if (compositor->log_file) {
        fclose(compositor->log_file);
        compositor->log_file = NULL;
    }
    
    if (compositor->scene_layout) {
        wlr_scene_output_layout_destroy(compositor->scene_layout);
    }
    struct wl_listener *listener, *tmp;
    wl_list_for_each_safe(listener, tmp, &compositor->keyboards, link) {
        wl_list_remove(&listener->link);
        free(listener);
    }
    wl_list_for_each_safe(listener, tmp, &compositor->pointers, link) {
        wl_list_remove(&listener->link);
        free(listener);
    }
}

static void handle_new_xdg_surface(struct wl_listener *listener, void *data) {
    struct onatu_compositor *compositor = wl_container_of(listener, compositor, new_xdg_surface);
    struct wlr_xdg_surface *xdg_surface = data;

    if (xdg_surface->role != WLR_XDG_SURFACE_ROLE_TOPLEVEL) return;

    struct onatu_window *window = calloc(1, sizeof(struct onatu_window));
    window->xdg_surface = xdg_surface;

    window->scene_node = wlr_scene_surface_create(
        &compositor->scene->tree,
        xdg_surface->surface
    );
    if (!window->scene_node) {
        LOG(compositor, "ERROR: Failed to create scene node for new window");
        free(window);
        return;
    }
    
    struct wlr_scene_rect *shadow = wlr_scene_rect_create(
        compositor->scene,
        10, 10, 0x404040FF);
    wlr_scene_node_lower_to_bottom(&shadow->node);

    wlr_scene_node_set_position(
        window->scene_node,
        compositor->outputs.next ? 0 : 200, 
        compositor->outputs.next ? 0 : 200
    );
    window->is_focused = false;
    window->is_maximized = false;
    window->is_fullscreen = false;
    window->z_index = 0;
    wl_list_insert(&compositor->windows, &window->link);
    wlr_scene_node_raise_to_top(window->scene_node);

    window->destroy.notify = handle_window_destroy;
    wl_signal_add(&xdg_surface->events.destroy, &window->destroy);

    window->map.notify = handle_window_map;
    wl_signal_add(&xdg_surface->surface->events.map, &window->map);

    window->commit.notify = handle_window_commit;
    wl_signal_add(&xdg_surface->surface->events.commit, &window->commit);

    printf("New window created: %p\n", window);

    window->last_commit_time = time(NULL);
    window->render_stats.frames_dropped = 0;
    window->render_stats.average_frametime = 0;
}

static void handle_window_destroy(struct wl_listener *listener, void *data) {
    (void)data;
    struct onatu_window *window = wl_container_of(listener, window, destroy);
    wl_list_remove(&window->link);
    wl_list_remove(&window->destroy.link);
    wl_list_remove(&window->map.link);
    wl_list_remove(&window->commit.link);
    printf("Window destroyed: %p\n", window);
    wlr_scene_node_destroy(window->scene_node);
    free(window->animation);
    free(window);
}

static void handle_window_map(struct wl_listener *listener, void *data) {
    (void)data;
    struct onatu_window *window = wl_container_of(listener, window, map);
 
    animation_init(&window->animation, 200.0f, EASE_LINEAR);
    window->animation.target_x = window->x;
    window->animation.target_y = window->y;
    window->animation.target_width = window->width;
    window->animation.target_height = window->height;
    window->animation.update_callback = update_window_position;
    window->animation.user_data = window;
    
    printf("Window mapped: %p\n", window);
}

static void handle_window_commit(struct wl_listener *listener, void *data) {
    (void)data;
    struct onatu_window *window = wl_container_of(listener, window, commit);

    if (animation_update(&window->animation) == 1) {
        wlr_scene_node_set_position(&window->scene_node->node, 
            window->animation.target_x, window->animation.target_y);
        wlr_scene_subsurface_tree_set_size(window->xdg_surface->surface,
            window->animation.target_width, window->animation.target_height);
    }
    
    printf("Window committed: %p\n", window);
}

void render_frame() {
    CHECK_GL_ERROR(glClear(GL_COLOR_BUFFER_BIT));
    
    if (!validate_shader_programs()) {
        log_error("Shader program validation failed");
        fallback_to_software_rendering();
        return;
    }
    
    struct onatu_output *output, *tmp_output;
    wl_list_for_each_safe(output, tmp_output, &compositor->outputs, link) {
        wl_list_remove(&output->link);
        free(output);
    }

    wlr_backend_destroy(compositor->backend);
    wl_display_destroy(compositor->display);
}

static void output_configure(struct wlr_output *output, float scale) {
    struct wlr_output_state state;
    wlr_output_state_init(&state);
    wlr_output_state_set_scale(&state, scale);
    wlr_output_commit_state(output, &state);
    wlr_output_state_finish(&state);
}

static void handle_new_output(struct wl_listener *listener, void *data) {
    struct onatu_compositor *compositor = wl_container_of(listener, compositor, new_output);
    struct wlr_output *wlr_output = data;

    struct onatu_output *output = calloc(1, sizeof(struct onatu_output));
    if (!output) {
        fprintf(stderr, "Failed to allocate output\n");
        return;
    }

    output->scene_output = wlr_scene_output_create(compositor->scene, wlr_output);
    wlr_scene_output_layout_add_output(compositor->scene_layout, output->scene_output, output->wlr_output);

    int next_x = 0;
    struct onatu_output *last_output;
    wl_list_for_each(last_output, &compositor->outputs, link) {
        next_x += last_output->wlr_output->width;
    }

    struct wlr_output_layout_output *layout_output = wlr_output_layout_add(compositor->scene_layout, wlr_output, next_x, 0);

    output_configure(wlr_output, wlr_output->scale);

    output_init(output, wlr_output, compositor->allocator, compositor->renderer, compositor);
    output->compositor = compositor;

    wl_list_insert(&compositor->outputs, &output->link);
    wlr_scene_output_commit(output->scene_output, NULL);
}

static void handle_new_input(struct wl_listener *listener, void *data) {
    struct onatu_compositor *compositor = wl_container_of(listener, compositor, new_input);
    struct wlr_input_device *device = data;

    switch (device->type) {
        case WLR_INPUT_DEVICE_KEYBOARD:
            keyboard_add(compositor, device);
            break;
        case WLR_INPUT_DEVICE_POINTER:
            pointer_add(compositor, device);
            break;
        default:
            break;
    }
}

void focus_window(struct onatu_compositor *compositor, struct onatu_window *window) {
    if (window && window->xdg_surface && !wlr_surface_is_xdg_surface(window->xdg_surface->surface)) return;
    if (compositor->focused_window == window) return;

    if (compositor->focused_window) {
        compositor->focused_window->is_focused = false;
    }

    compositor->focused_window = window;
    window->is_focused = true;

    struct wlr_keyboard *keyboard = wlr_seat_get_keyboard(compositor->seat);
    if (keyboard) {
        wlr_seat_keyboard_notify_enter(compositor->seat, window->xdg_surface->surface,
                                      keyboard->keycodes, keyboard->num_keycodes, &keyboard->modifiers);
    }
}

void compositor_move_window(struct onatu_window *window, int x, int y) {
    window->x = x;
    window->y = y;
    wlr_scene_node_set_position(&window->scene_node->node, x, y);
}

void compositor_set_title(struct onatu_window *window, const char *title) {
    wlr_xdg_toplevel_set_title(window->xdg_surface->toplevel, title);
    free(window->title);
    window->title = strdup(title);
}

void tile_windows(struct onatu_compositor *compositor) {
    if (wl_list_empty(&compositor->windows)) return;

    struct onatu_output *output;
    wl_list_for_each(output, &compositor->outputs, link) {
        int width = output->wlr_output->width;
        int height = output->wlr_output->height;
        
        int window_count = 0;
        struct onatu_window *w;
        wl_list_for_each(w, &compositor->windows, link) window_count++;

        int cols = (int)ceil(sqrt(window_count));
        int rows = (window_count + cols - 1) / cols;
        
        int cell_width = width / cols;
        int cell_height = height / rows;
        
        int i = 0;
        wl_list_for_each(w, &compositor->windows, link) {
            int col = i % cols;
            int row = i / cols;
            
            w->x = output->x + col * cell_width + 5;
            w->y = output->y + row * cell_height + 5;
            wlr_xdg_toplevel_set_size(w->xdg_surface, 
                cell_width - 10, cell_height - 10);
            
            if (i == 0) focus_window(compositor, w);
            i++;
        }
    }
}

void stack_windows(struct onatu_compositor *compositor) {
    struct onatu_window *w;
    int y_offset = 50;
    wl_list_for_each(w, &compositor->windows, link) {
        w->x = 50;
        w->y = y_offset;
        wlr_xdg_toplevel_set_size(w->xdg_surface, 600, 400);
        y_offset += 450;
    }
}

void sort_windows_by_zindex(struct onatu_compositor *compositor) {
    struct onatu_window *windows[64];
    int count = 0;

    struct onatu_window *win;
    wl_list_for_each(win, &compositor->windows, link) {
        if (count < 64) windows[count++] = win;
    }

    for (int i = 0; i < count-1; i++) {
        for (int j = 0; j < count-i-1; j++) {
            if (windows[j]->z_index < windows[j+1]->z_index) {
                struct onatu_window *temp = windows[j];
                windows[j] = windows[j+1];
                windows[j+1] = temp;
            }
        }
    }

    wl_list_remove(&compositor->windows);
    wl_list_init(&compositor->windows);
    for (int i = 0; i < count; i++) {
        wl_list_insert(&compositor->windows, &windows[i]->link);
    }
}

void float_window(struct onatu_window *window) {
    window->is_floating = true;

    struct wlr_box geometry = {
        .x = 100,
        .y = 100,
        .width = 400,
        .height = 300,
    };
    wlr_xdg_toplevel_set_size(window->xdg_surface->toplevel, geometry.width, geometry.height);
    wlr_xdg_surface_schedule_configure(window->xdg_surface);
}

void maximize_window(struct onatu_compositor *compositor, struct onatu_window *window) {
    if (window->animation.active) return;
    
    composinator_animate_window(window, 0, 0, compositor->max_width, compositor->max_height);
    window->is_maximized = true;
    wlr_xdg_toplevel_set_size(window->xdg_surface->toplevel,
        compositor->max_width,
        compositor->max_height);
}

void unmaximize_window(struct onatu_window *window) {
    if (!window->is_maximized) return;

    struct wlr_xdg_toplevel *toplevel = window->xdg_surface->toplevel;
    wlr_xdg_toplevel_set_maximized(toplevel, false);
    window->is_maximized = false;
}

void LOG(struct onatu_compositor *comp, const char *fmt, ...) {
    va_list args;
    va_start(args, fmt);
    vfprintf(comp->log_file, fmt, args);
    fprintf(comp->log_file, "\n");
    va_end(args);
}

void keyboard_add(struct onatu_compositor *compositor, struct wlr_input_device *device) {
    struct wlr_keyboard *keyboard = wlr_keyboard_from_input_device(device);
    
    struct wl_listener *keyboard_key = calloc(1, sizeof(struct wl_listener));
    keyboard_key->notify = handle_keyboard_key;
    wl_signal_add(&keyboard->events.key, keyboard_key);

    struct gesture_context *ctx = calloc(1, sizeof(struct gesture_context));
    ctx->compositor = compositor;
    wlr_keyboard_set_user_data(keyboard, ctx);
}

static void handle_pointer_swipe_begin(struct wl_listener *listener, void *data) {
    struct gesture_context *ctx = wl_container_of(listener, ctx, swipe_begin);
    ctx->gesture_data.swipe.active = true;
    ctx->gesture_data.swipe.fingers = ((struct wlr_pointer_swipe_begin_event*)data)->fingers;
}

static void handle_pointer_pinch_update(struct wl_listener *listener, void *data) {
    struct gesture_context *ctx = wlr_container_of(listener, ctx, pinch_update);
    struct wlr_pointer_pinch_update_event *event = data;
    
    if(ctx->gesture_data.pinch.scale != event->scale) {
        handle_zoom_gesture(ctx->compositor, event->scale);
    }
}

float animation_interpolate(enum animation_ease ease, float t) {
    switch(ease) {
    case EASE_IN_OUT_CUBIC:
        return t < 0.5 ? 4*t*t*t : 1-pow(-2*t+2,3)/2;
    case EASE_OUT_BACK: {
        float c1 = 1.70158;
        return 1 + c1*pow(t-1,3) + c1*pow(t-1,2);
    }
    default:
        return t;
    }
}

int animation_update(struct animation *anim) {
    const float stiffness = 400.0f;
    const float damping = 25.0f;
    const float dt = 0.016f;

    float dx = anim->target_x - anim->x;
    float ax = stiffness * dx - damping * anim->velocity_x;
    anim->velocity_x += ax * dt;
    anim->x += anim->velocity_x * dt;

    float dy = anim->target_y - anim->y;
    float ay = stiffness * dy - damping * anim->velocity_y;
    anim->velocity_y += ay * dt;
    anim->y += anim->velocity_y * dt;

    return fabsf(dx) > 0.1f || fabsf(dy) > 0.1f;
}
